import argparse
from typing import Dict, Any, Optional
import toml

from utils.toolkit import NamespaceDict


def load_config(cfg_path) -> Dict[str, Any]:
    """
    加载TOML配置文件并处理变量引用
    支持$variable形式的变量引用
    """
    # 读取原始配置
    config = toml.load(cfg_path)

    # 处理变量引用
    process_variable_references(config)

    return config


def merge_config(cfg, args: Optional[argparse.Namespace]) -> Dict[str, Any]:
    """merge argparse args to config, replace the value in config if in args"""
    # keys = config.keys()
    if args is not None:
        for key, value in vars(args).items():
            if value is not None:
                cfg[key] = value
    return cfg


def cfg_dealer(cfg_path, args: Optional[argparse.Namespace]) -> NamespaceDict:
    """
    处理配置字典，合并命令行参数
    """
    # 合并命令行参数到配置字典
    config: Dict[str, Any] = load_config(cfg_path)
    config = merge_config(config, args)
    process_variable_references(config)
    cfg = NamespaceDict(**config)

    # postprocess dataset manager config
    cfg.dataset.update(cfg.general)
    # postprocess training config
    cfg.training.general.inc_epochs += [cfg.training.general.inc_epochs[-1]] * (
        cfg.general.nb_tasks - len(cfg.training.general.inc_epochs)
    )
    cfg.training.general.epochs = [
        cfg.training.general.init_epochs
    ] + cfg.training.general.inc_epochs

    cfg.training.general.inc_lr += [cfg.training.general.inc_lr[-1]] * (
        cfg.general.nb_tasks - len(cfg.training.general.inc_lr)
    )
    cfg.training.optimizer.lr = [
        cfg.training.general.init_lr
    ] + cfg.training.general.inc_lr
    cfg.training.scheduler.T_max = [
        cfg.training.general.init_epochs
    ] + cfg.training.general.inc_epochs

    return NamespaceDict(**cfg)


"""Deal with variable references ($Var) in the config file"""


def process_variable_references(config_dict, root_dict=None):
    """
    递归处理配置字典中的变量引用
    变量引用格式为 $variable_name
    """
    if root_dict is None:
        root_dict = config_dict

    for key, value in list(config_dict.items()):
        if isinstance(value, dict):
            process_variable_references(value, root_dict)
        elif isinstance(value, str) and value.startswith("$"):
            var_name = value[1:]  # 去掉$前缀

            # 在整个配置中查找变量值
            var_value = find_variable_value(var_name, root_dict)

            if var_value is not None:
                config_dict[key] = var_value
            else:
                print(f"Warning: Variable {var_name} referenced but not found")


def find_variable_value(var_path, config_dict):
    """
    在配置字典中查找指定路径的变量值
    支持点表示法，例如 training.general.init_epochs
    """
    if "." in var_path:
        parts = var_path.split(".")
        current = config_dict
        for part in parts:
            if part in current:
                current = current[part]
            else:
                return None
        return current
    else:
        # 扁平搜索所有嵌套字典
        return flat_search(var_path, config_dict)


def flat_search(key, nested_dict):
    """
    在嵌套字典中平面搜索指定的键
    """
    if key in nested_dict:
        return nested_dict[key]

    for k, v in nested_dict.items():
        if isinstance(v, dict):
            result = flat_search(key, v)
            if result is not None:
                return result

    return None
